package com.xs.tool.formula;

import lombok.extern.slf4j.Slf4j;

import java.util.*;
import java.util.concurrent.ThreadLocalRandom;

/// 抽卡计算公式 k:返回值 v:权重 - 使用别名方法优化到O(1)抽卡
@Slf4j
public class DrawCardsFormula<T> extends HashMap<T, Long> {
    public final Long weightPool;

    // 别名方法所需的数据结构
    private final T[] items;
    private final double[] probabilities;  // 每个位置的概率
    private final int[] aliases;          // 别名数组
    private final int size;


    public static <T> DrawCardsFormula<T> fromIntegerMap(Map<T, Integer> map) {
        return new DrawCardsFormula<>(convertIntegerMapToLongMap(map));
    }

    // 辅助方法：将Map<T, Integer>转换为Map<T, Long>
    @SuppressWarnings("unused")
    private static <T> Map<T, Long> convertIntegerMapToLongMap(Map<T, Integer> map) {
        Map<T, Long> longMap = new HashMap<>();
        for (Map.Entry<T, Integer> entry : map.entrySet()) {
            longMap.put(entry.getKey(), entry.getValue().longValue());
        }
        return longMap;
    }

    @SuppressWarnings("unchecked")
    public DrawCardsFormula(Map<T, Long> map) {
        super(map);

        // 过滤掉权重为0的项目
        List<Map.Entry<T, Long>> validEntries = map.entrySet().stream()
                .filter(entry -> entry.getValue() > 0)
                .toList();

        this.size = validEntries.size();
        if (size == 0) {
            throw new IllegalArgumentException("没有有效的权重项目");
        }

        weightPool = validEntries.stream().mapToLong(Map.Entry::getValue).sum();

        // 初始化别名方法数据结构
        items = (T[]) new Object[size];
        probabilities = new double[size];
        aliases = new int[size];

        // 构建别名表
        buildAliasTable(validEntries);

        log.debug("权重抽卡公式初始化完成(别名方法)，总权重池: {}, 有效物品数量: {}", weightPool, size);
    }

    /**
     * 构建别名表 - O(n) 预处理
     */
    private void buildAliasTable(List<Map.Entry<T, Long>> entries) {
        double[] normalizedProbs = new double[size];
        double avgProb = 1.0 / size;

        // 归一化概率
        for (int i = 0; i < size; i++) {
            items[i] = entries.get(i).getKey();
            normalizedProbs[i] = (double) entries.get(i).getValue() / weightPool * size;
        }

        // 分离小概率和大概率
        Queue<Integer> small = new ArrayDeque<>();
        Queue<Integer> large = new ArrayDeque<>();

        for (int i = 0; i < size; i++) {
            if (normalizedProbs[i] < 1.0) {
                small.offer(i);
            } else {
                large.offer(i);
            }
        }

        // 构建别名表
        while (!small.isEmpty() && !large.isEmpty()) {
            int less = small.poll();
            int more = large.poll();

            probabilities[less] = normalizedProbs[less];
            aliases[less] = more;

            normalizedProbs[more] = (normalizedProbs[more] + normalizedProbs[less]) - 1.0;

            if (normalizedProbs[more] < 1.0) {
                small.offer(more);
            } else {
                large.offer(more);
            }
        }

        // 处理剩余项目
        while (!large.isEmpty()) {
            probabilities[large.poll()] = 1.0;
        }
        while (!small.isEmpty()) {
            probabilities[small.poll()] = 1.0;
        }
    }

    /// 权重抽卡 - 批量抽取
    public List<T> draw(int num) {
        List<T> results = new ArrayList<>(num);

        for (int i = 0; i < num; i++) {
            T drawnItem = drawSingle();
            if (drawnItem != null) {
                results.add(drawnItem);
            }
        }

        return results;
    }

    /// 单次权重抽卡 - O(1) 时间复杂度
    public T drawSingle() {
        if (size == 0) {
            log.warn("权重抽卡失败，物品列表为空");
            return null;
        }

        ThreadLocalRandom random = ThreadLocalRandom.current();

        // O(1) 别名方法抽卡
        int index = random.nextInt(size);
        double prob = random.nextDouble();

        if (prob < probabilities[index]) {
            return items[index];
        } else {
            return items[aliases[index]];
        }
    }

    /// 特殊优化：权重相等时的快速抽卡
    public static <T> DrawCardsFormula<T> createUniform(List<T> items) {
        Map<T, Long> uniformWeights = new HashMap<>();
        for (T item : items) {
            uniformWeights.put(item, 1L);
        }
        return new DrawCardsFormula<>(uniformWeights);
    }

    /**
     * 传统二分查找实现（用于性能对比）
     */
    public static class BinarySearchFormula<T> {
        private final List<T> items;
        private final long[] cumulativeWeights;
        private final long weightPool;

        public BinarySearchFormula(Map<T, Long> map) {
            items = new ArrayList<>();
            List<Long> weights = new ArrayList<>();

            for (Entry<T, Long> entry : map.entrySet()) {
                if (entry.getValue() > 0) {
                    items.add(entry.getKey());
                    weights.add(entry.getValue());
                }
            }

            cumulativeWeights = new long[items.size()];
            long cumulative = 0;
            for (int i = 0; i < weights.size(); i++) {
                cumulative += weights.get(i);
                cumulativeWeights[i] = cumulative;
            }
            weightPool = cumulative;
        }

        public T drawSingle() {
            if (weightPool <= 0) return null;

            long randomValue = ThreadLocalRandom.current().nextLong(weightPool);
            int left = 0, right = items.size() - 1;

            while (left < right) {
                int mid = left + (right - left) / 2;
                if (cumulativeWeights[mid] <= randomValue) {
                    left = mid + 1;
                } else {
                    right = mid;
                }
            }

            return items.get(left);
        }
    }

    public static void performanceTest(Map<String, Long> weights, int drawCount) {
        log.info("开始性能测试，物品数量: {}, 抽卡次数: {}", weights.size(), drawCount);

        // 测试别名方法
        long startTime = System.nanoTime();
        DrawCardsFormula<String> aliasFormula = new DrawCardsFormula<>(weights);
        long constructTime = System.nanoTime() - startTime;

        startTime = System.nanoTime();
        for (int i = 0; i < drawCount; i++) {
            aliasFormula.drawSingle();
        }
        long aliasDrawTime = System.nanoTime() - startTime;

        // 测试传统二分查找方法
        startTime = System.nanoTime();
        DrawCardsFormula.BinarySearchFormula<String> binaryFormula = new DrawCardsFormula.BinarySearchFormula<>(weights);
        long binaryConstructTime = System.nanoTime() - startTime;

        startTime = System.nanoTime();
        for (int i = 0; i < drawCount; i++) {
            binaryFormula.drawSingle();
        }
        long binaryDrawTime = System.nanoTime() - startTime;

        // 输出结果
        log.info("=== 性能测试结果 ===");
        log.info("别名方法 - 构造时间: {} ms, 抽卡时间: {} ms", constructTime / 1_000_000.0, aliasDrawTime / 1_000_000.0);
        log.info("二分查找 - 构造时间: {} ms, 抽卡时间: {} ms", binaryConstructTime / 1_000_000.0, binaryDrawTime / 1_000_000.0);
        log.info("抽卡性能提升: {}x", (double) binaryDrawTime / aliasDrawTime);
    }
}
